package org.argus.job.server.system.allocate;

import cn.hutool.core.collection.CollectionUtil;
import com.alibaba.nacos.api.exception.NacosException;
import com.alibaba.nacos.api.naming.pojo.Instance;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import lombok.extern.slf4j.Slf4j;
import org.argus.common.core.utils.MathUtils;
import org.argus.common.core.utils.StringUtils;
import org.argus.common.redis.service.RedisService;
import org.argus.job.common.model.AppBindingInfo;
import org.argus.job.common.model.JobServerInstance;
import org.argus.job.server.system.domain.AppInfo;
import org.argus.job.server.system.management.mapper.AppInfoMapper;
import org.redisson.api.RLock;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.Collectors;

import static org.argus.job.server.common.constants.JobCacheConstants.APP_ASSIGNED_TOPIC;
import static org.argus.job.server.common.constants.JobCacheConstants.APP_SERVER_ASSIGN_KEY;

/**
 * App 分配
 * 
 * @Author 刘文/Cliff Liu
 * @Date 2024/7/25
 */
@Component
@Slf4j
public class AppAllocator {

    @Autowired
    private JobServerInstanceSubscriber jobServerInstanceSubscriber;

    @Autowired
    private RedissonClient redissonClient;

    @Autowired
    private AppInfoMapper appInfoMapper;

    @Autowired
    private RedisService redisService;

    @Autowired
    private RedisTemplate redisTemplate;

    private static final String APP_ALLOCATE_KEY_NAME = "job:app:allocate";

    public void allocate() {
        List<Instance> runningServers;
        try {
            runningServers = jobServerInstanceSubscriber.getAllInstances();
        } catch (NacosException e) {
            log.error("Get job server instances from nacos failed.", e);
            return;
        }
        // 新增App，需强制分配，避免其它Server正在分配但没感知新应用，导致新App无法分配到服务器
        allocate(runningServers, true);
    }

    /**
     * 根据正在运行的Job Server实例列表分配App
     * 
     * @param runningServers
     * @param force 是否强制分配
     */
    public void allocate(List<Instance> runningServers, boolean force) {
        if (CollectionUtil.isEmpty(runningServers)) {
            log.info("No running servers, keep everything the same.");
            return;
        }
        RLock lock = redissonClient.getLock(APP_ALLOCATE_KEY_NAME);
        try {
            boolean locked = false;
            while (!locked) {
                locked = lock.tryLock(100, 1000, TimeUnit.MILLISECONDS);
                // 非强制分配，没获取到锁，说明其它节点正在分配，本节点分配不执行
                if (!force && !locked) {
                    log.warn("Return for another server is allocating the apps");
                    return;
                }
            }
            try {
                doAllocation(runningServers);
                Thread.sleep(200);
            } catch (Exception e) {
                log.error("Allocate app failed", e);
            } finally {
                lock.unlock();
            }
        } catch (InterruptedException e) {
            log.error("Get lock failed", e);
        }
    }

    private void doAllocation(List<Instance> runningServers) {
        Map<JobServerInstance, LinkedList<AppInfo>> runningServerMap =
            runningServers.stream().map(instance -> new JobServerInstance(instance.getIp(), instance.getPort()))
                .collect(Collectors.toMap(Function.identity(), o -> new LinkedList<>()));
        List<AppInfo> appInfos = appInfoMapper.getNormalApps();
        List<AppInfo> downServerApps = new LinkedList<>();

        // APP CODE -> Server IP 分配结果
        List<AppBindingInfo> appBindingInfos = new ArrayList<>(appInfos.size());
        for (AppInfo appInfo : appInfos) {
            if (StringUtils.isBlank(appInfo.getCurrentServer())) {
                downServerApps.add(appInfo);
                continue;
            }
            JobServerInstance bindJobServer = JobServerInstance.of(appInfo.getCurrentServer());
            LinkedList<AppInfo> serverApps = runningServerMap.get(bindJobServer);
            // App所属Server正在运行
            if (serverApps != null) {
                serverApps.add(appInfo);
                // 正在运行Server的App不重新分配
                AppBindingInfo bindingInfo = new AppBindingInfo();
                bindingInfo.setAppId(appInfo.getId().toString());
                bindingInfo.setAppName(appInfo.getName());
                bindingInfo.setInstance(bindJobServer);
                appBindingInfos.add(bindingInfo);
            } else {
                downServerApps.add(appInfo);
            }
        }
        allocateDownServerApps(downServerApps, appInfos, runningServerMap, appBindingInfos);
        // 将分配结果写入Redis缓存
        redisService.setCacheObject(APP_SERVER_ASSIGN_KEY, appBindingInfos);
        // 推送App已分配事件
        redisTemplate.convertAndSend(APP_ASSIGNED_TOPIC, appInfos.get(0));
        log.info("App allocation is saved to redis");
    }

    private void allocateDownServerApps(List<AppInfo> downServerApps, List<AppInfo> appInfos,
        Map<JobServerInstance, LinkedList<AppInfo>> runningServerMap, List<AppBindingInfo> appBindingInfos) {
        // 没有因服务器下线而需要分配的应用
        if (CollectionUtil.isEmpty(downServerApps)) {
            log.info("No app need to allocate");
            return;
        }

        log.info("Have {} apps to re-allocate.", downServerApps.size());
        int averageAppCount = MathUtils.ceilDivide(appInfos.size(), runningServerMap.size());
        Map<JobServerInstance, List<AppInfo>> newAllocatedServerApps =
            runningServerMap.keySet().stream().collect(Collectors.toMap(Function.identity(), o -> new LinkedList<>()));
        for (AppInfo appInfo : downServerApps) {
            for (Map.Entry<JobServerInstance, List<AppInfo>> entry : newAllocatedServerApps.entrySet()) {
                // 当前服务器新分配app + 原来管理中的app数量 < 平均分配数量
                if (entry.getValue().size() + runningServerMap.get(entry.getKey()).size() < averageAppCount) {
                    entry.getValue().add(appInfo);
                    AppBindingInfo bindingInfo = new AppBindingInfo();
                    bindingInfo.setAppId(appInfo.getId().toString());
                    bindingInfo.setAppName(appInfo.getName());
                    bindingInfo.setInstance(entry.getKey());
                    appBindingInfos.add(bindingInfo);
                    break;
                }
            }
        }

        AtomicInteger appCount = new AtomicInteger();
        // 更新数据库中app所属服务器信息
        newAllocatedServerApps.forEach((server, apps) -> {
            for (AppInfo appInfo : apps) {
                LambdaUpdateWrapper<AppInfo> updateWrapper = Wrappers.lambdaUpdate(AppInfo.class)
                    .set(AppInfo::getCurrentServer, server.toString()).eq(AppInfo::getId, appInfo.getId());
                appInfoMapper.update(null, updateWrapper);
                appCount.incrementAndGet();
            }
        });
        log.info("Allocated {} apps", appCount.get());
    }

}
